import tensorflow as tf
from tensorflow.keras import layers, models
import numpy as np
import matplotlib.pyplot as plt
import os
import datetime

# 初始化TensorBoard可用性标志
TENSORBOARD_AVAILABLE = False

# 尝试导入TensorBoard，如果失败则跳过
try:
    from tensorflow.keras.callbacks import TensorBoard
    TENSORBOARD_AVAILABLE = True
except ImportError:
    print("TensorBoard不可用，将跳过相关功能")

# 设置GPU内存增长，避免一次性占用所有GPU内存
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    try:
        # 设置GPU内存增长
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
        print(f"检测到 {len(gpus)} 个GPU设备，已设置内存增长模式")
    except RuntimeError as e:
        print(f"设置GPU内存增长失败: {e}")
else:
    print("未检测到GPU设备，将使用CPU进行计算")

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

def load_local_mnist(data_path="./mnist_data"):
    """
    从本地目录加载MNIST数据
    """
    print(f"尝试从本地目录 {data_path} 加载MNIST数据...")
    
    # 检查是否存在本地MNIST数据文件
    mnist_npz_path = os.path.join(data_path, "mnist.npz") if os.path.isdir(data_path) else data_path
    
    if os.path.exists(mnist_npz_path) and mnist_npz_path.endswith('.npz'):
        try:
            # 从本地NPZ文件加载数据
            with np.load(mnist_npz_path, allow_pickle=True) as f:
                x_train, y_train = f['x_train'], f['y_train']
                x_test, y_test = f['x_test'], f['y_test']
            print(f"成功从本地文件 {mnist_npz_path} 加载数据")
            return (x_train, y_train), (x_test, y_test)
        except Exception as e:
            print(f"从本地NPZ文件加载数据失败: {e}")
    
    # 如果本地加载失败，使用tf.keras方式加载
    print("使用TensorFlow内置方式加载MNIST数据...")
    return tf.keras.datasets.mnist.load_data()

def preprocess_image(image):
    """预处理单张图像"""
    # 扩展维度以适应图像处理函数
    image = tf.expand_dims(image, axis=-1)
    # 调整图像大小以适应简化模型 (32x32)
    image = tf.image.resize(image, [32, 32])
    # 转换为RGB
    image = tf.image.grayscale_to_rgb(image)
    # 归一化
    image = tf.cast(image, tf.float32) / 255.0
    return image

def load_and_preprocess_data():
    """加载并预处理MNIST数据"""
    print("正在加载MNIST数据集...")
    (x_train, y_train), (x_test, y_test) = load_local_mnist("../MNIST_data")
    
    # 转换为float32以节省内存
    x_train = x_train.astype(np.float32)
    x_test = x_test.astype(np.float32)
    
    # 标签one-hot编码
    y_train = tf.keras.utils.to_categorical(y_train, 10)
    y_test = tf.keras.utils.to_categorical(y_test, 10)
    
    print(f"训练集形状: {x_train.shape}")
    print(f"测试集形状: {x_test.shape}")
    print(f"训练标签形状: {y_train.shape}")
    print(f"测试标签形状: {y_test.shape}")
    
    return (x_train, y_train), (x_test, y_test)

def create_simple_model():
    """创建简化版CNN模型"""
    model = models.Sequential([
        # 卷积层1
        layers.Conv2D(32, (3, 3), activation='relu', input_shape=(32, 32, 3)),
        layers.MaxPooling2D((2, 2)),
        
        # 卷积层2
        layers.Conv2D(64, (3, 3), activation='relu'),
        layers.MaxPooling2D((2, 2)),
        
        # 卷积层3
        layers.Conv2D(64, (3, 3), activation='relu'),
        
        # 展平层
        layers.Flatten(),
        
        # 全连接层
        layers.Dense(64, activation='relu'),
        layers.Dropout(0.5),
        
        # 输出层
        layers.Dense(10, activation='softmax')
    ])
    
    return model

def visualize_samples(x_data, y_data, num_samples=9):
    """可视化数据样本"""
    plt.figure(figsize=(8, 8))
    for i in range(num_samples):
        plt.subplot(3, 3, i+1)
        # 预处理图像用于显示
        processed_img = preprocess_image(x_data[i])
        plt.imshow(processed_img[:,:,0], cmap='gray')
        plt.title(f'标签: {np.argmax(y_data[i])}')
        plt.axis('off')
    plt.tight_layout()
    plt.show()

def plot_training_history(history):
    """绘制训练历史"""
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
    
    # 绘制损失曲线
    ax1.plot(history.history['loss'], label='训练损失')
    ax1.plot(history.history['val_loss'], label='验证损失')
    ax1.set_title('模型损失')
    ax1.set_xlabel('轮次')
    ax1.set_ylabel('损失')
    ax1.legend()
    ax1.grid(True)
    
    # 绘制准确率曲线
    ax2.plot(history.history['accuracy'], label='训练准确率')
    ax2.plot(history.history['val_accuracy'], label='验证准确率')
    ax2.set_title('模型准确率')
    ax2.set_xlabel('轮次')
    ax2.set_ylabel('准确率')
    ax2.legend()
    ax2.grid(True)
    
    plt.tight_layout()
    plt.show()

# 创建自定义数据生成器以节省内存
class MNISTDataGenerator(tf.keras.utils.Sequence):
    def __init__(self, x_data, y_data, batch_size=32, shuffle=True):
        self.x_data = x_data
        self.y_data = y_data
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.indices = np.arange(len(x_data))
        self.on_epoch_end()
    
    def __len__(self):
        return int(np.ceil(len(self.x_data) / self.batch_size))
    
    def __getitem__(self, index):
        # 获取当前批次的索引
        start_idx = index * self.batch_size
        end_idx = min((index + 1) * self.batch_size, len(self.x_data))
        batch_indices = self.indices[start_idx:end_idx]
        
        # 处理当前批次的数据
        batch_x = np.zeros((len(batch_indices), 32, 32, 3), dtype=np.float32)
        batch_y = self.y_data[batch_indices]
        
        for i, idx in enumerate(batch_indices):
            # 预处理单张图像
            batch_x[i] = preprocess_image(self.x_data[idx]).numpy()
        
        return batch_x, batch_y
    
    def on_epoch_end(self):
        if self.shuffle:
            np.random.shuffle(self.indices)

def main():
    """主函数"""
    print("MNIST手写数字识别 - 简化版CNN实现")
    print("=" * 50)
    
    # 加载数据
    (x_train, y_train), (x_test, y_test) = load_and_preprocess_data()
    
    # 可视化样本数据
    print("\n可视化样本数据...")
    visualize_samples(x_train, y_train)
    
    # 创建数据生成器以节省内存
    print("\n创建数据生成器...")
    train_generator = MNISTDataGenerator(x_train, y_train, batch_size=32, shuffle=True)
    test_generator = MNISTDataGenerator(x_test, y_test, batch_size=32, shuffle=False)
    
    # 创建模型
    print("\n创建简化版CNN模型...")
    model = create_simple_model()
    
    # 显示模型结构
    model.summary()
    
    # 编译模型
    model.compile(
        optimizer='adam',
        loss='categorical_crossentropy',
        metrics=['accuracy']
    )
    
    # 设置TensorBoard回调（如果可用）
    callbacks = []
    tensorboard_enabled = TENSORBOARD_AVAILABLE  # 使用局部变量避免全局变量修改问题
    log_dir = None
    
    if tensorboard_enabled:
        try:
            # 使用简单的英文路径，避免中文字符问题
            base_log_dir = "../logs_simple"
            # 确保基础日志目录存在
            os.makedirs(base_log_dir, exist_ok=True)
            # 创建基于时间戳的子目录
            timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
            log_dir = os.path.join(base_log_dir, timestamp)
            # 创建完整的日志目录
            os.makedirs(log_dir, exist_ok=True)
            tensorboard_callback = TensorBoard(log_dir=log_dir, histogram_freq=1)
            callbacks.append(tensorboard_callback)
            print(f"TensorBoard功能已启用，日志将保存到: {log_dir}")
        except Exception as e:
            print(f"创建TensorBoard日志目录失败: {e}")
            print("TensorBoard功能将被禁用")
            tensorboard_enabled = False
    
    # 训练模型
    print("\n开始训练模型...")
    try:
        history = model.fit(
            train_generator,
            epochs=5,
            verbose=1,
            validation_data=test_generator,
            callbacks=callbacks,  # 添加TensorBoard回调（如果可用）
            # 添加这些参数以减少警告
            workers=1,
            use_multiprocessing=False
        )
    except Exception as e:
        print(f"训练过程中出现错误: {e}")
        print("尝试使用CPU进行训练...")
        # 禁用GPU，使用CPU训练
        os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
        # 重新编译模型
        model.compile(
            optimizer='adam',
            loss='categorical_crossentropy',
            metrics=['accuracy']
        )
        # 使用更小的批次大小再次尝试
        train_generator = MNISTDataGenerator(x_train, y_train, batch_size=16, shuffle=True)
        test_generator = MNISTDataGenerator(x_test, y_test, batch_size=16, shuffle=False)
        
        history = model.fit(
            train_generator,
            epochs=5,
            verbose=1,
            validation_data=test_generator,
            callbacks=callbacks,  # 添加TensorBoard回调（如果可用）
            # 添加这些参数以减少警告
            workers=1,
            use_multiprocessing=False
        )
    
    # 绘制训练历史
    print("\n绘制训练历史...")
    plot_training_history(history)
    
    # 评估模型
    print("\n评估模型...")
    test_loss, test_accuracy = model.evaluate(test_generator, verbose=0)
    print(f"测试损失: {test_loss:.4f}")
    print(f"测试准确率: {test_accuracy:.4f}")
    
    # 预测示例
    print("\n预测示例...")
    # 获取前10个测试样本进行预测
    sample_x = np.zeros((10, 32, 32, 3), dtype=np.float32)
    for i in range(10):
        sample_x[i] = preprocess_image(x_test[i]).numpy()
    
    predictions = model.predict(sample_x)
    predicted_classes = np.argmax(predictions, axis=1)
    true_classes = np.argmax(y_test[:10], axis=1)
    
    print("前10个测试样本的预测结果:")
    for i in range(10):
        print(f"样本 {i+1}: 真实标签={true_classes[i]}, 预测标签={predicted_classes[i]}, "
              f"置信度={np.max(predictions[i]):.4f}")
    
    # 保存模型
    model.save('simple_cnn_model.h5')
    print("\n模型已保存为 'simple_cnn_model.h5'")
    
    # 提供TensorBoard启动命令
    if tensorboard_enabled and callbacks:
        print("\n训练完成！要查看TensorBoard，请在终端中运行以下命令：")
        print(f"tensorboard --logdir {log_dir}")

if __name__ == "__main__":
    main()